import torch
import pickle
import math
import random
from torchvision import transforms
from Cifar10ConvNet import ConvNet
import numpy as np
import matplotlib.pyplot as plt
import torchvision

# CIFAR10数据集的输出是范围在[0,1]之间的PILImage，我们将它转换并归一化范围在[-1,1]之间的Tensor
transforms_test = transforms.Compose([
    # ToTensor会将数据除以255数据[0, 1]，把HWC变成CHW
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

test_dataset = torchvision.datasets.CIFAR10(root='./resources/',
                                            train=False,
                                            transform=transforms_test)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=15,
                                          shuffle=False)

classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

model = ConvNet()
model.eval()
model.load_state_dict(torch.load('./ckpt/Cifar10Model.ckpt'))


# encoding='iso-8859-1'
def unpickle(file):
    with open(file, 'rb') as f:
        dict = pickle.load(f, encoding='bytes')
    return dict


def imshow(img, label):
    """
    # 反归一
    img = img * 0.5 + 0.5
    npimg = np.transpose(img.numpy(), (1, 2, 0))
    plt.imshow(npimg)
    plt.show()
    """
    row = math.ceil(len(img)/5)
    for i in range(len(img)):
        output = model(img[i].reshape(1, 3, 32, 32))
        predicted = output.argmax(dim=1)
        plt.subplot(row, 5, i+1)
        image = img[i]*0.5+0.5
        image = np.transpose(image.numpy(), (1, 2, 0))
        plt.imshow(image)
        plt.axis("off")
        plt.title("{}, ({})".format(classes[predicted], classes[label[i]]),
                  color="green" if classes[label[i]] == classes[predicted] else "red")
    plt.show()


test_file = "./resources/cifar-10-batches-py/test_batch"
dictData = unpickle(test_file)
data = dictData[b'data']
labels = dictData[b'labels']

with torch.no_grad():

    testData = iter(test_loader)
    image, label = testData.__next__()

    # 将若干副图拼成一副图像
    # imshow(torchvision.utils.make_grid(image))
    imshow(image, label)

    correct = 0
    for i in np.arange(15):
        img = np.reshape(data[i], (3, 32, 32))
        # 将图片转化成张量
        # transform = transforms.ToTensor()
        images = transforms_test(img.reshape(32, 32, 3))
        # images = transform(img)
        images = images.reshape(1, 3, 32, 32)
        output = model(images)
        # _, predicted = torch.max(output.data, 1)
        predicted = output.argmax(dim=1)
        if labels[i] == predicted:
            correct += 1

        img = img.transpose(1, 2, 0)
        plt.subplot(3, 5, i + 1)
        plt.axis("off")
        plt.imshow(img)
        plt.title("{},({})".format(classes[int(predicted)], classes[labels[i]]),
                  color="green" if labels[i] == int(predicted) else "red")

    print(correct/15)
    plt.show()
    '''
    showNum = 1
    rows = math.ceil(showNum / 5)
    for i in range(showNum):
        index = random.randint(0, 1)
        img = np.reshape(data[index], (3, 32, 32))
        # 将图片转化成张量
        # transform = transforms.ToTensor()
        images = transforms_test(img.reshape(32, 32, 3))
        images = images.reshape(1, 3, 32, 32)
        output = model(images)
        predict = output.argmax(dim=1)
        # plt.imshow只能接受形如(32，32，3)
        img = img.transpose(1, 2, 0)
        plt.subplot(rows, 5, i + 1)
        plt.axis("off")
        plt.imshow(img)
        plt.title("{},({})".format(classes[int(predict)], classes[labels[index]]),
                  color="green" if labels[index] == int(predict) else "red")
    plt.show()

    '''



